/**
 * @file Naiver_stokes.cpp
 * @author LSJ (you@domain.com)
 * @brief 
 * @version 0.1
 * @date 2021-01-18
 * 
 * @copyright Copyright (c) 2021
 * 
 */
#include "Mesh.h"
#include "Element.h"
#include "MultigridSolver.h"
#include <Eigen/Sparse>
#include <cmath>
#include "MINRES.h"
#include "FEMFunction.h"

#define pi 4 * atan(1.0)
#define vis 0.01
typedef Eigen::SparseMatrix<Real> SpMat;
typedef Eigen::Triplet<Real> Tri;
typedef Eigen::VectorXd VectorXd;

/// 边界条件
Real bnd_ux(Real* p)
{
    Real x = p[0];
    Real y = p[1];
    //return 1 - y * y;
    return 2 * x * y * y * y;
}

Real bnd_uy(Real* p)
{
    Real x = p[0];
    Real y = p[1];
    //return 0;
	return  0.5* pow(x,4) - 0.5 * pow(y, 4);
}

double bnd_p(const double * p)
{
    return 0;
}

double f(double * p)
{
    return 0;
}

int main(int argc, char* argv[])
{
    int n = 1;
    RectangleDomain * Domain = new RectangleDomain({{-1,-1}, {1,-1}, {1,1}, {-1,1}});
    Mesh<2> * meshU = new Q2Mesh(Domain,{POW2(n), POW2(n)});
    Mesh<2> * meshP = new Q1Mesh(Domain,{POW2(n), POW2(n)});
    Element<2> * Uelement = new Quadrilateral_2_Element();
    Element<2> * Pelement = new Quadrilateral_1_Element();

    int n_u = meshU->n_dofs();
    int n_p = meshP->n_dofs();
    int total_n_dof = 2 * n_u + n_p;
    int n_UeleDof = Uelement -> n_Dofs();
    int n_PeleDof = Pelement -> n_Dofs();
    /// 组装刚度矩阵
    SpMat StiffMat(total_n_dof, total_n_dof);
	SpMat A;
	SpMat Q(n_p,n_p);
    std::vector<Tri> TriList(meshU->n_element() * n_UeleDof * (2*n_UeleDof + 4*n_PeleDof));
	std::vector<Tri> TriListQ(meshP->n_element() * n_PeleDof *n_PeleDof);
    std::vector<Tri>::iterator it = TriList.begin();
	std::vector<Tri>::iterator itq = TriListQ.begin();
    for (int k = 0; k < meshU->n_element(); k++)
    {
	std::vector<int> Idx = meshU->NodeofEle(k);
	std::vector<Dofs<2> > temp(n_UeleDof);
	for (int i = 0; i < n_UeleDof; i++)
	    temp[i] = meshU->DofsofIndex(Idx[i]);
	Uelement->SetDofsList(temp);
	Idx = meshP->NodeofEle(k);
	temp.resize(n_PeleDof);
	for (int i = 0; i < n_PeleDof; i++)
	    temp[i] = meshP->DofsofIndex(Idx[i]);
	Pelement->SetDofsList(temp);

	for (int i = 1; i <= n_UeleDof; i++)
	{
	    for (int j = 1; j <= n_UeleDof; j++)
	    {
		Real a = 0;
		for (int gp = 0; gp < Uelement->n_GaussPnt(); gp++)
		{
		    Real xi = Uelement->GaussionPoint(gp)[0];
		    Real eta = Uelement->GaussionPoint(gp)[1];
		    a += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * InnerProuduct(Uelement->gradient(xi, eta, i), Uelement->gradient(xi, eta, j));
		}
		*it = Tri(Uelement->NdIdx(i), Uelement->NdIdx(j), a);
		it++;
		*it = Tri(Uelement->NdIdx(i)+n_u, Uelement->NdIdx(j)+n_u, a);
		it++;
	    }
	    for (int j = 1; j <= n_PeleDof; j++)
	    {
		Real bx = 0;
		Real by = 0;
		for (int gp = 0; gp < Uelement->n_GaussPnt(); gp++)
		{
		    Real xi = Uelement->GaussionPoint(gp)[0];
		    Real eta = Uelement->GaussionPoint(gp)[1];
		    bx += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * Pelement->phi(xi, eta, j) * Uelement->phi_x(xi, eta, i);
		    by += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * Pelement->phi(xi, eta, j) * Uelement->phi_y(xi, eta, i);
		}
		*it = Tri(Pelement->NdIdx(j)+2*n_u, Uelement->NdIdx(i), -bx);
		it++;
		*it = Tri(Uelement->NdIdx(i), Pelement->NdIdx(j)+2*n_u, -bx);
		it++;
		*it = Tri(Pelement->NdIdx(j)+2*n_u, Uelement->NdIdx(i)+n_u, -by);
		it++;
		*it = Tri(Uelement->NdIdx(i)+n_u, Pelement->NdIdx(j)+2*n_u, -by);
		it++;		
	    }
	}	
		for(int i = 1;i <= n_PeleDof;i++)
		{
			for(int j = 1;j <= n_PeleDof;j++)
			{
				Real q = 0;
				for (int gp = 0; gp < Pelement->n_GaussPnt(); gp++)
				{
		    	Real xi = Pelement->GaussionPoint(gp)[0];
		    	Real eta = Pelement->GaussionPoint(gp)[1];
		    	q += Pelement->det_Jacobi(xi, eta) * Pelement->GaussionWeight(gp) * Pelement->phi(xi,eta,i) * Pelement->phi(xi,eta,j);
				}
				*itq = Tri(Pelement->NdIdx(i),Pelement->NdIdx(j),q);
				itq++;
			}
		}
    }
    StiffMat.setFromTriplets(TriList.begin(), TriList.end());
    StiffMat.makeCompressed();
    Q.setFromTriplets(TriListQ.begin(),TriListQ.end());
	Q.makeCompressed();
	
    /// 处理边界条件
    VectorXd rhs = Eigen::MatrixXd::Zero(total_n_dof, 1);
    for (auto k:meshU->Boundary())
    {
	int bnd_Idxx = k;
	int bnd_Idxy = k + n_u;
	Dofs<2> bnd_point = meshU -> DofsofIndex(k);
	Real bnd_valuex = bnd_ux(*bnd_point);
	Real bnd_valuey = bnd_uy(*bnd_point);
	rhs[bnd_Idxx] = bnd_valuex * StiffMat.coeffRef(bnd_Idxx, bnd_Idxx);
	for (Eigen::SparseMatrix<Real>::InnerIterator it(StiffMat, bnd_Idxx);it;++it)
	{
	    int row = it.row();
	    if (row == bnd_Idxx)
		continue;
	    StiffMat.coeffRef(bnd_Idxx, row) = 0.0;
	    rhs[row] -= StiffMat.coeffRef(row, bnd_Idxx) * bnd_valuex;
	    StiffMat.coeffRef(row, bnd_Idxx) = 0.0;	    
	}
	for (Eigen::SparseMatrix<Real>::InnerIterator it(StiffMat, bnd_Idxy);it;++it)
	{
	    int row = it.row();
	    if (row == bnd_Idxy)
		continue;
	    StiffMat.coeffRef(bnd_Idxy, row) = 0.0;
	    rhs[row] -= StiffMat.coeffRef(row, bnd_Idxy) * bnd_valuey;
	    StiffMat.coeffRef(row, bnd_Idxy) = 0.0;	    
	}
    }
	A = StiffMat.block(0,0,n_u,n_u);
	StokesPreconditioner<double> stokespred(A,Q);
	Eigen::MinRes<SpMat,StokesPreconditioner<double>> solver(stokespred);
    solver.compute(StiffMat);
    VectorXd x = solver.solve(rhs);
	std::cout << " stokes equaion complete." << std::endl;
    FEMFunction<2> ux(x.segment(0,n_u));
    FEMFunction<2> uy(x.segment(n_u,n_u));
    FEMFunction<2> pp(x.tail(n_p));
    //stokes的解作为Navier_stokes方程牛顿迭代
    SpMat F(total_n_dof,total_n_dof);
    VectorXd rhss(total_n_dof);
    VectorXd solution = x;
    VectorXd delta(total_n_dof);
    std::vector<Tri> TriListF(meshU->n_element() * n_UeleDof * (8*n_UeleDof + 4*n_PeleDof));
    std::vector<Tri>::iterator itf = TriListF.begin();
    int step = 0;
    double error = 1e-12;
    double res = 1;
    double h = 0.1;
    while (res > error)
    {
        for (int k = 0; k < meshU->n_element(); k++)
        {
	    std::vector<int> Idx = meshU->NodeofEle(k);
	    std::vector<Dofs<2> > temp(n_UeleDof);
	    for (int i = 0; i < n_UeleDof; i++)
	        temp[i] = meshU->DofsofIndex(Idx[i]);
	    Uelement->SetDofsList(temp);
	    Idx = meshP->NodeofEle(k);
	    temp.resize(n_PeleDof);
	    for (int i = 0; i < n_PeleDof; i++)
	        temp[i] = meshP->DofsofIndex(Idx[i]);
	    Pelement->SetDofsList(temp);
        for (int i = 1; i <= n_UeleDof; i++)
	    {
	    for (int j = 1; j <= n_UeleDof; j++)
	    {
            long int ind_i = Uelement->NdIdx(i);
            long int ind_j = Uelement->NdIdx(j);
            double a = 0.0;
            double n = 0.0;
            double wxx = 0.0;
            double wyy = 0.0;
            double wxy = 0.0;
            double wyx = 0.0;
            for (int gp = 0; gp < Uelement->n_GaussPnt(); gp++)
		    {

		    Real xi = Uelement->GaussionPoint(gp)[0];
		    Real eta = Uelement->GaussionPoint(gp)[1];
            Real Jxy = Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp);
		    a += Jxy * InnerProuduct(Uelement->gradient(xi, eta, i), Uelement->gradient(xi, eta, j));
		    n += Jxy * (ux.value({xi,eta},Uelement) * Uelement->gradient(xi,eta,j)[0] + uy.value({xi,eta},Uelement) * Uelement->gradient(xi,eta,j)[1])*Uelement->phi(xi,eta,i);
            wxx += Jxy * (ux.gradient({xi,eta},Uelement)[0])* Uelement->phi(xi,eta,i) * Uelement->phi(xi,eta,j);
            wyy += Jxy * (uy.gradient({xi,eta},Uelement)[1])* Uelement->phi(xi,eta,i) * Uelement->phi(xi,eta,j);
            wxy += Jxy * (ux.gradient({xi,eta},Uelement)[1])* Uelement->phi(xi,eta,i) * Uelement->phi(xi,eta,j);
            wyx += Jxy * (uy.gradient({xi,eta},Uelement)[0])* Uelement->phi(xi,eta,i) * Uelement->phi(xi,eta,j);
            }
            *itf = Tri(ind_i,ind_j,a);
            itf++;
            *itf = Tri(ind_i + n_u,ind_j + n_u,a);
            itf++;
            *itf = Tri(ind_i,ind_j,n);
            itf++;
            *itf = Tri(ind_i + n_u,ind_j + n_u,n);
            itf++;
            *itf = Tri(ind_i,ind_j,wxx);
            itf++;
            *itf = Tri(ind_i + n_u,ind_j + n_u,wxx);
            itf++;
            *itf = Tri(ind_i,ind_j + n_u,wxy);
            itf++;
            *itf = Tri(ind_i + n_u,ind_j,wyx);
            itf++;
        }
        for (int j = 1; j <= n_PeleDof; j++)
	    {
		    Real bx = 0;
		    Real by = 0;
		    for (int gp = 0; gp < Uelement->n_GaussPnt(); gp++)
		    {
		    Real xi = Uelement->GaussionPoint(gp)[0];
		    Real eta = Uelement->GaussionPoint(gp)[1];
		    bx += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * Pelement->phi(xi, eta, j) * Uelement->phi_x(xi, eta, i);
		    by += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * Pelement->phi(xi, eta, j) * Uelement->phi_y(xi, eta, i);
		    }
		    *it = Tri(Pelement->NdIdx(j)+2*n_u, Uelement->NdIdx(i), -bx);
		    it++;
		    *it = Tri(Uelement->NdIdx(i), Pelement->NdIdx(j)+2*n_u, -bx);
		    it++;
		    *it = Tri(Pelement->NdIdx(j)+2*n_u, Uelement->NdIdx(i)+n_u, -by);
		    it++;
		    *it = Tri(Uelement->NdIdx(i)+n_u, Pelement->NdIdx(j)+2*n_u, -by);
		    it++;		
	    }
        }
        }
        F.setFromTriplets(TriListF.begin(),TriListF.end());
        F.makeCompressed();
        for (int k = 0; k < meshU->n_element(); k++)
        {
	    std::vector<int> Idx = meshU->NodeofEle(k);
	    std::vector<Dofs<2> > temp(n_UeleDof);
	    for (int i = 0; i < n_UeleDof; i++)
	        temp[i] = meshU->DofsofIndex(Idx[i]);
	    Uelement->SetDofsList(temp);
	    Idx = meshP->NodeofEle(k);
	    temp.resize(n_PeleDof);
	    for (int i = 0; i < n_PeleDof; i++)
	        temp[i] = meshP->DofsofIndex(Idx[i]);
	    Pelement->SetDofsList(temp);
        for(int gp = 0;gp < Uelement->n_GaussPnt();gp++)
        {
            Real xi = Uelement->GaussionPoint(gp)[0];
		    Real eta = Uelement->GaussionPoint(gp)[1];
            double Jxy = Uelement -> det_Jacobi(xi,eta) * Uelement ->GaussionWeight(gp);
            for(int i = 1;i <= n_UeleDof;i++)
            {
                double cont1 = Jxy *((-ux.value({xi,eta},Uelement)* ux.gradient({xi,eta},Uelement)[0]-uy.value({xi,eta},Uelement)* ux.gradient({xi,eta},Uelement)[1])) * Uelement->phi(xi,eta,i);
                cont1 -= Jxy *(vis * InnerProuduct(ux.gradient({xi,eta},Uelement),Uelement->gradient(xi,eta,i)));
                cont1 += Jxy * pp.value({xi,eta},Pelement) * Uelement->gradient(xi,eta,i)[0];
                rhs(Uelement->NdIdx(i)) += cont1;
                double cont2 = Jxy *((-ux.value({xi,eta},Uelement)* uy.gradient({xi,eta},Uelement)[0]-uy.value({xi,eta},Uelement)* uy.gradient({xi,eta},Uelement)[1]));
                cont2 -= Jxy *(vis * InnerProuduct(uy.gradient({xi,eta},Uelement),Uelement->gradient(xi,eta,i)));
                cont2 += Jxy * pp.value({xi,eta},Pelement) * Uelement->gradient(xi,eta,i)[1];
                rhss(Uelement->NdIdx(i) + n_u) += cont1;
            }
            for(int i = 1;i <= n_PeleDof;i++)
            {
                double cont = Jxy * Pelement->phi(xi,eta,i) * (ux.gradient({xi,eta},Uelement)[0] + ux.gradient({xi,eta},Uelement)[0]);
                rhss(Pelement->NdIdx(i) + 2 * n_u) += cont;
            }
        }
        }//end element loop;

        for (auto k:meshU->Boundary())
        {
	    int bnd_Idxx = k;
	    int bnd_Idxy = k + n_u;
	    Dofs<2> bnd_point = meshU -> DofsofIndex(k);
	    Real bnd_valuex = bnd_ux(*bnd_point);
	    Real bnd_valuey = bnd_uy(*bnd_point);
	    rhs[bnd_Idxx] = bnd_valuex * F.coeffRef(bnd_Idxx, bnd_Idxx);
	    for (Eigen::SparseMatrix<Real>::InnerIterator it(F, bnd_Idxx);it;++it)
	    {
	        int row = it.row();
	        if (row == bnd_Idxx)
		    continue;
	        F.coeffRef(bnd_Idxx, row) = 0.0;
	        rhss[row] -= StiffMat.coeffRef(row, bnd_Idxx) * bnd_valuex;
	        F.coeffRef(row, bnd_Idxx) = 0.0;	    
	    }
	    for (Eigen::SparseMatrix<Real>::InnerIterator it(F, bnd_Idxy);it;++it)
	    {
	        int row = it.row();
	        if (row == bnd_Idxy)
		    continue;
	        F.coeffRef(bnd_Idxy, row) = 0.0;
	        rhss[row] -= StiffMat.coeffRef(row, bnd_Idxy) * bnd_valuey;
	        F.coeffRef(row, bnd_Idxy) = 0.0;	    
	    }
        }
		std::cout << "F is :" << F << std::endl;
		std::cout << "rhs is :" << rhss << std::endl;
        Eigen::BiCGSTAB<SpMat,Eigen::IncompleteLUT<double> > Solver;
        Solver.compute(F);
        Solver.setTolerance(1e-12);
        std::cout << "Newton iteration step begin..." << std::endl;
        delta = Solver.solve(rhss);
        std::cout << "Newton iteration step end..." << std::endl;
        for (int i = 0; i < n_u; i++)
	        ux(i) += delta(i);
	    for (int i = n_u; i < 2 * n_u; i++)
	        uy(i - n_u) += delta(i);
	    for (int i = 2 * n_u; i < total_n_dof; i++)
	        pp(i - 2 * n_u) += delta(i);

        for (int i = 0; i < n_u; i++)
	        solution(i) = ux(i);
	    for (int i = n_u; i < 2 * n_u; i++)
	        solution(i) = uy(i - n_u);
	    for (int i = 2 * n_u; i < total_n_dof; i++)
	        solution(i) = pp(i - 2 * n_u);

        res = sqrt(delta.squaredNorm());
        step++;
        std::cout << "nolinear res updated at "<<step << " th step is : " << res << std::endl;
    }//end while
}